package com.cwj.algorithm.graph;

import com.cwj.algorithm.priority.IndexPriorityQueue;
import com.cwj.algorithm.uf.UnionFindDisjointSet;

import java.util.Deque;
import java.util.LinkedList;

/**
 * @author chenwujie
 * @date 2020-12-25 14:05
 */
public class KruskalMST {
    private EdgeWeightedGraph edgeWeightedGraph;
    /**
     * 保存最小生成树的边
     */
    private Deque<Edge> mst;
    /**
     * 已关联到最小生成树的顶点
     */
    private UnionFindDisjointSet uf;
    /**
     * 最小索引优先队列维护待处理的边
     */
    private IndexPriorityQueue<Edge> pq;

    public KruskalMST(EdgeWeightedGraph edgeWeightedGraph){
        this.edgeWeightedGraph = edgeWeightedGraph;
        this.uf = new UnionFindDisjointSet(edgeWeightedGraph.getV());
        this.pq = new IndexPriorityQueue<>(edgeWeightedGraph.getE(), true);
        this.mst = new LinkedList<>();
    }

    private void mst(){
        // 遍历所有的边添加到最小索引优先队列
        for (Edge edge : edgeWeightedGraph.getEdges()) {
            pq.insert(edge);
        }

        // 遍历队列，把边对应的顶点添加到最小生成树中
        while(mst.size() < (edgeWeightedGraph.getV() - 1) && !pq.isEmpty()){
            int i = pq.removeFirst();
            Edge edge = edgeWeightedGraph.getEdges().get(i);
            int v = edge.getV();
            int w = edge.getW();

            if(uf.connected(v, w)){
                // 如果不在一颗树上
                continue;
            }
            // 合并到一颗树
            uf.union(v, w);
            // 该边为最小生成树的一条边
            mst.push(edge);
        }
    }

    public Deque<Edge> createMst(){
        mst();
        return mst;
    }
}
